import os
import shutil
from PIL import Image, ImageStat
from statistics import median

def calculate_std_dev(image_path):
    img = Image.open(image_path)
    stat = ImageStat.Stat(img)
    std_dev = sum(stat.stddev) / 3
    return std_dev

def find_overexposed_images(folder_path, z_exposure_path, std_threshold):
    for root, dirs, files in os.walk(folder_path):
        for file in files:
            file_path = os.path.join(root, file)
            if file.endswith('.jpg') or file.endswith('.jpeg') or file.endswith('.png'):
                std_dev = calculate_std_dev(file_path)
                if std_dev < std_threshold:
                    source_folder = os.path.dirname(file_path)
                    target_folder = os.path.join(z_exposure_path, source_folder.split(os.sep)[-1])
                    if not os.path.exists(target_folder):
                        os.makedirs(target_folder)
                    target_file_path = os.path.join(target_folder, file)
                    shutil.move(file_path, target_file_path)
                    print(f"Moved {file_path} to {target_folder}")

def determine_std_threshold(sample_folder):
    std_devs = []
    for root, dirs, files in os.walk(sample_folder):
        for file in files:
            file_path = os.path.join(root, file)
            if file.endswith('.jpg') or file.endswith('.jpeg') or file.endswith('.png'):
                std_dev = calculate_std_dev(file_path)
                std_devs.append(std_dev)
                # print(f"Image: {file_path}, Standard Deviation: {std_dev}")
    
    # 假设样本图片中包含正常曝光和过度曝光的图片,选择最小的标准差作为初始阈值
    initial_threshold = median(std_devs) * 2
    print(f"Initial threshold: {initial_threshold}")
    
    return initial_threshold

# 设置要查找的文件夹路径
folder_path = 'traindata_ori_bak_20240320'

# 设置要移动的图片的文件夹路径
z_exposure_path = 'z_exposure'

# 设置包含样本图片的文件夹路径
sample_folder = 'exposure_samples'

# 确定标准差阈值
std_threshold = determine_std_threshold(sample_folder)

# 根据阈值查找并移动完全曝光过度的图片
find_overexposed_images(folder_path, z_exposure_path, std_threshold)